#include "ProcessInjectorImpl_Win.h"
#include "tools/buffer.h"
#include "tools/helpMacros.h"
#include "tools/impl/win/WinExpection.h"
#include <boost/lexical_cast.hpp> 
#include "Psapi.h"
#include <Winternl.h>
#include SD_TL_PLATFORM_ARH_INC( .,ProcessInjector )
#include SD_TL_PLATFORM_IMPL_INC( tools,HandleWraper )

using namespace sd;
using namespace prin;
using namespace impl;
using namespace tools;
using namespace buffer;





bool ProcessImpl::isStarted()const
{
	return !hThread_;
}


void ProcessImpl::start()
{
	assert(hThread_!=0);
	ResumeThread(hThread_);
	hThread_=0;
}



/*
	asm

	InterlockedCompareExchange

	
	push pointer_to_lib
	LoadLibrary



*/



LPTHREAD_START_ROUTINE ProcessImpl::getJumper()
{
	if(!jumper_)
	{
		jumper_ = static_cast<LPTHREAD_START_ROUTINE>(getProcAddr("LoadLibraryW"));
		/*
		const int AllocSize=sizeof(inject_stub);
		void* jmp = VirtualAllocEx( hProcess_,NULL, AllocSize,  MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE );

		if (!jumper_)
		{
			//TODO Expection
		}

		inject_stub injector(jmp);

		DWORD lSize;
		if( !WriteProcessMemory( hProcess_, jumper_, &injector, AllocSize, &lSize ) || sizeof(injector)!=lSize )
		{
			//TODO Expection
		}

		jumper_=static_cast<LPTHREAD_START_ROUTINE>(jmp);
		*/
	}
	return jumper_;
}
#include <iostream>

void ProcessImpl::Inject( const std::wstring& name, const std::wstring& dllParams)
{
	const int AllocSize=(name.size()+1)*sizeof(std::wstring::value_type);;

	//TODO function for write data as buffer
	void* addr = VirtualAllocEx( hProcess_,NULL, AllocSize,  MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE );
	//TODO Expection check
	SIZE_T lSize;
	if( !WriteProcessMemory( hProcess_, addr, name.c_str(), AllocSize, &lSize ) || AllocSize!=lSize )
	{
		//TODO Expection
	}
	
	DWORD threadId;
	HANDLE thread = CreateRemoteThread(hProcess_, NULL, 0, getJumper(), addr, dllParams.size()?CREATE_SUSPENDED:0, &threadId );

	if(!thread)
	{
		//TODO Expection
	}

	if(dllParams.size())
		pushCallParams(threadId,dllParams);

	CloseHandle(thread);
}


void* ProcessImpl::getLoadLibraryAddr()
{
	HMODULE lib = LoadLibraryA("Kernel32.dll");

	void* addr = GetProcAddress( lib, "LoadLibraryW" );

	CloseHandle( lib );

	return addr;
}



std::wstring ProcessImpl::getCallParamsName(int threadid, int processid)
{
	std::wstringstream s;
	s<< L"HookMainParams_"<<processid<< L"_" << threadid;
	return s.str();
}

void ProcessImpl::pushCallParams(int threadid, const std::wstring& params)
{
	std::wstring fname = getCallParamsName( threadid, getId());
	const size_t buf_size = (params.size()+1)*sizeof(std::wstring::value_type);

	HHANDLE  hMapFile = HHANDLE( CreateFileMappingW( INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, buf_size, fname.c_str()));

	if(!hMapFile)
	{
//TODO expection
	}

	wchar_t *pBuf = (wchar_t *) MapViewOfFile(hMapFile.impl(), FILE_MAP_ALL_ACCESS, 0, 0, buf_size);

	if(!pBuf)
	{
//TODO expection
	}

	memcpy( pBuf, params.c_str(), buf_size);
    UnmapViewOfFile(pBuf); 
}



std::wstring ProcessImpl::loadName()const
{
	wchar_t buffer[ 1024];
	if(!GetModuleFileNameExW ( hProcess_, NULL, &buffer[0], 1024*sizeof(wchar_t)))
		throw win::WinImplException(std::wstring(L"Error load process for ")+boost::lexical_cast<std::wstring>(getId()) );
	return buffer;
}






typedef decltype(NtQueryInformationProcess)* TNtQueryInformationProcess;


std::wstring ProcessImpl::loadCmdName()const
{

	HMODULE lib = 0;
	std::wstring result;

	try
	{
		lib = LoadLibraryW(L"Ntdll.dll");
		if(!lib)
			throw;

		TNtQueryInformationProcess procedure = (TNtQueryInformationProcess)(GetProcAddress(lib, "NtQueryInformationProcess" ));
		if(!procedure)
			throw;

		PROCESS_BASIC_INFORMATION pbi;
		ULONG pbis;
		if(0!=(*procedure)(hProcess_, ProcessBasicInformation, &pbi, sizeof(pbi), &pbis))
			throw;

		result = (*pbi.PebBaseAddress).ProcessParameters->CommandLine.Buffer;
	}catch(...)
	{
		wchar_t buffer[MAX_PATH];
		GetProcessImageFileNameW(hProcess_, buffer, MAX_PATH*2);
		result = buffer;
	}
	return result;
}



std::wstring ProcessImpl::getName()const
{
	if(name_.empty())
		name_ = loadName();
	return name_;
}


int ProcessImpl::getId()const
{
	return GetProcessId( hProcess_ );
}



std::wstring ProcessImpl::getCmd()const
{
	if(cmdName_.empty())
		cmdName_ = loadCmdName();
	return cmdName_;
}

ProcessImpl::ProcessImpl( HANDLE hProcess, HANDLE hThread, bool started, const std::wstring& cmdName ): hProcess_(hProcess), hThread_(started?NULL:hThread), cmdName_(cmdName),jumper_(NULL){}




ProcessInjectorImpl::ProcessInjectorImpl():privExist_(false),IsWow64Process_( static_cast<TIsWow64Process>( getProcAddr("IsWow64Process"))){}



bool ProcessInjectorImpl::checkX6compatibility( HANDLE process )
{
	bool result = true;

	if( IsWow64Process_ )
	{
		BOOL currentIs64;
		IsWow64Process( GetCurrentProcess(), &currentIs64 );
		if(currentIs64)
		{
			BOOL processIs64;
			IsWow64Process( process, &processIs64 );
			result = !!processIs64;
		}
	}
	return result;
}


ProcessInjectorImpl::ProcessImplRef ProcessInjectorImpl::createProcess(const std::wstring& name, bool start)
{
	DynamicBuffer cmdBuff( PtrBuffer( name.c_str(), ( name.length() + 1 ) * sizeof(std::wstring::value_type) ) );
	
	STARTUPINFOW startupInfo;
	memset( &startupInfo, 0, sizeof(startupInfo));
	startupInfo.cb = sizeof(startupInfo);

	PROCESS_INFORMATION processInformation;
	
	BOOL  result = CreateProcessW( NULL, &(cmdBuff.cast<wchar_t>()), NULL, NULL, FALSE, CREATE_PRESERVE_CODE_AUTHZ_LEVEL | (!start?CREATE_SUSPENDED:0), NULL, NULL, &startupInfo, &processInformation );


	if(!checkX6compatibility(processInformation.hProcess))
		throw win::WinImplException(std::wstring(L"Using x32 loader for x64 process (")+name+L") not allowed");


	if(!result)
		throw win::WinImplException(std::wstring(L"Error when create process ")+name);

	return ProcessInjectorImpl::ProcessImplRef(new ProcessImpl( processInformation.hProcess, processInformation.hThread, start, name ));
};




ProcessInjectorImpl::ProcessImplRef ProcessInjectorImpl::attachToProcess( int uid )
{
	checkPrivileges();

	HANDLE hProcess = OpenProcess( PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ | PROCESS_SET_INFORMATION, FALSE, uid );

	if( !hProcess )
		throw win::WinImplException(std::wstring(L"Error when attaching  to process with handle ")+boost::lexical_cast<std::wstring>(uid));

	if(!checkX6compatibility(hProcess))
		throw win::WinImplException(std::wstring(L"Using x32 loader for x64 process not allowed"));

	return ProcessInjectorImpl::ProcessImplRef(new ProcessImpl( hProcess ));
}

void SetPrivilege(
    HANDLE hProcess,          // access token handle
    LPCTSTR lpszPrivilege,  // name of privilege to enable/disable
    BOOL bEnablePrivilege   // to enable or disable privilege
    ) 
{

	HANDLE hToken;
	if(!OpenProcessToken( hProcess, TOKEN_ADJUST_PRIVILEGES, &hToken ))
	{
		return;
	}



    TOKEN_PRIVILEGES tp;
    LUID luid;

    if ( !LookupPrivilegeValue( 
            NULL,            // lookup privilege on local system
            lpszPrivilege,   // privilege to lookup 
            &luid ) )        // receives LUID of privilege
    {
        return ; 
    }

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    if (bEnablePrivilege)
        tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    else
        tp.Privileges[0].Attributes = 0;

    // Enable the privilege or disable all privileges.

    if ( !AdjustTokenPrivileges(
           hToken, 
           FALSE, 
           &tp, 
           sizeof(TOKEN_PRIVILEGES), 
           (PTOKEN_PRIVILEGES) NULL, 
           (PDWORD) NULL) )
    { 
          return ; 
    } 

    if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
    {

          return ;
    } 
}





void ProcessInjectorImpl::grantPrivileges()
{
	SetPrivilege( GetCurrentProcess(), SE_DEBUG_NAME, true );
}

